from concurrent.futures import ThreadPoolExecutor

import requests
from Crypto.Cipher import AES
import os
import subprocess
import concurrent.futures
import shutil
import re

headers = {
    'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36'
}


class M3u8Searcher:
    def __init__(self):
        pass

    # @staticmethod
    def search(self, search):
        # todo
        return 'video_url', 'aes_key', 'aes_key_url', 'iv'


class AnotherSearcher(M3u8Searcher):
    def __init__(self):
        super().__init__()

    def search(self, search):
        # todo
        return 'video_url', 'aes_key', 'aes_key_url', 'iv'


class M3u8Source:
    """

    """

    def __init__(self, search, video_url, aes_key, aes_key_url, iv, **kwargs):
        """

        """
        self.search = search
        self.video_url = video_url
        self.aes_key = aes_key
        self.aes_key_url = aes_key_url
        self.iv = iv
        if not video_url:
            self.get_source(kwargs.get('searcher') or M3u8Searcher())

    def get_source(self, searcher: M3u8Searcher):
        self.video_url, self.aes_key, self.aes_key_url, self.iv = searcher.search(self.search)


class M3u8Download:
    _source = None
    _save_dir = None

    def __init__(self, save_dir, source: M3u8Source):
        self._source = source
        self._save_dir = save_dir
        os.makedirs(save_dir, exist_ok=True)

    def download(self):
        if self._source:
            self.get_video_m3u8()
            self.merge_ts(self._save_dir, self._source.search)
        else:
            print('source not inited')

    def get_video_m3u8(self):
        """
        解析m3u8资源 下载ts列表并解密（如需要）
        ...
        :return:
        """
        try:
            save_dir, video_url, aes_key, aes_key_url, iv = self._save_dir, self._source.video_url, self._source.aes_key, self._source.aes_key_url, self._source.iv
            video_m3u8_file = os.path.join(save_dir, 'index.m3u8')
            if os.path.exists(video_m3u8_file):
                with open(video_m3u8_file, 'rb') as f:
                    video_m3u8_content = f.read()
            else:
                video_m3u8_content = requests.get(video_url, headers=headers).content
                with open(video_m3u8_file, 'wb') as f:
                    f.write(video_m3u8_content)
            video_m3u8_str = video_m3u8_content.decode('utf-8')
            if aes_key or aes_key_url:
                if not aes_key:
                    aes_key = requests.get(aes_key_url).content
            print(f'aes key {aes_key}')

            # 创建一个线程池，最大线程数为 12
            futures = []
            with concurrent.futures.ThreadPoolExecutor(max_workers=os.cpu_count() * 4) as executor:
                ts_lines = [l for l in re.split(r'\r\n?|\n', video_m3u8_str) if l.endswith('.ts')]
                self._ts_batch(video_url, save_dir, ts_lines, futures, executor, aes_key, iv)
                # 等待所有任务完成
                for future in concurrent.futures.as_completed(futures):
                    # 处理任务的结果
                    print(f"download complete {future.result()} of {len(ts_lines)}")
                    # 当with语句块结束时，线程池会自动关闭和清理
                completed_len = len(os.listdir(os.path.join(save_dir, 'decrypted')))
                if len(ts_lines) > completed_len:
                    print(f'ts list not complete all:{len(ts_lines)} completed: {completed_len}')
                    self._ts_batch(video_url, save_dir, ts_lines, futures, executor, aes_key, iv)
            os.remove(video_m3u8_file)
        except Exception as e:
            print(f'error {str(e)}')

    def _ts_batch(self, video_url, video_dir, ts_lines, futures, executor: ThreadPoolExecutor, aes_key, aes_iv):
        for i in range(0, len(ts_lines)):
            ts_source = ts_lines[i] if ts_lines[i].startswith('http') else video_url.rsplit('/', 1)[0] + '/' + ts_lines[
                i]
            ts_dir = os.path.join(video_dir, 'decrypted')
            if not os.path.exists(ts_dir):
                os.mkdir(ts_dir)
            ts_name = str(i) + '.ts'
            ts_file = os.path.join(ts_dir, ts_name)
            if not os.path.exists(ts_file) or os.stat(ts_file).st_size == 0:
                futures.append(executor.submit(self._ts_parse, video_dir, ts_name, ts_source, aes_key, aes_iv))

    @staticmethod
    def _ts_parse(video_dir, ts_name, ts_source, aes_key, iv):
        ts_dir = os.path.join(video_dir, 'decrypted')
        ts_file = os.path.join(ts_dir, ts_name)
        response = requests.get(ts_source, headers=headers)
        if response.status_code == 200:
            ts_encrypted = response.content
            with open(ts_file, 'wb') as f:
                # todo 拓展其他解密方式
                f.write(ts_encrypted if not aes_key else AES.new(aes_key, AES.MODE_CBC, iv).decrypt(ts_encrypted))
        else:
            print(f'state {response.status_code}')
        return ts_name

    @staticmethod
    def merge_ts(save_dir, save_name):
        ts_dir = os.path.join(save_dir, 'decrypted')
        os.makedirs(ts_dir, exist_ok=True)
        ts_names = [f for f in os.listdir(ts_dir) if f.endswith('.ts')]
        ts_orders = []
        for ts_name in ts_names:
            order = int(ts_name.split('.', 1)[0])
            ts_orders.append(order)
        file_list = os.path.join(save_dir, 'a_ts_list.txt')
        with open(file_list, 'w') as f:
            for ts_order in sorted(ts_orders):
                name = str(ts_order) + '.ts'
                f.write('file \'' + os.path.realpath(os.path.join(ts_dir, name)) + '\'\n')
        output_file = os.path.join(save_dir, f'{save_name}.mp4')
        commands = ['ffmpeg',
                    '-f', 'concat',
                    '-safe', '0',
                    '-i', file_list,
                    '-c', 'copy',
                    output_file]
        # 需要下载ffmpeg并配置path
        subprocess.run(commands, check=True)
        shutil.rmtree(ts_dir)
        os.remove(file_list)


if __name__ == '__main__':
    dest_dir = "D:\\AppData\\share\\jiongb"
    series_name = "囧爸"
    with open("media_source.txt", 'r') as media_source:
        series_source = re.split(r'\r\n?|\n', media_source.read())
    for i in range(len(series_source)):
        seg_name = f'{series_name}{i + 1}'
        if os.path.exists(os.path.join(dest_dir, f'{seg_name}.mp4 ')):
            print(f'existed {seg_name}')
            continue

        M3u8Download(dest_dir, M3u8Source(seg_name, series_source[i], None, None, None)).download()
